import torch
from torch import nn
from TorchCRF import CRF
from P03_NER.LSTM_CRF.utils.data_loader import *

config = Config()

class NERLSTM_CRF(nn.Module):
    def __init__(self, embedding_dim, hidden_dim, dropout, word2id, tag2id):
        super(NERLSTM_CRF, self).__init__()
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.vocab_size = len(word2id)
        self.tag_size = len(tag2id)

        self.name = 'BiLSTM_CRF'
        self.word_embeds = nn.Embedding(num_embeddings=self.vocab_size,
                                        embedding_dim=self.embedding_dim)
        self.dropout = nn.Dropout(p=dropout)
        self.lstm = nn.LSTM(input_size=self.embedding_dim,
                            hidden_size=self.hidden_dim // 2,
                            batch_first=True,
                            bidirectional=True)
        self.hidden2tag = nn.Linear(in_features=self.hidden_dim,
                                    out_features=self.tag_size)
        self.crf = CRF(self.tag_size)

    def forward(self, input, mask):
        outputs = self.get_lstm2linear(input)
        outputs = outputs * mask.unsqueeze(-1)
        outputs = self.crf.viterbi_decode(outputs, mask)
        return outputs

    def log_likelihood(self, x, tags, mask):
        outputs = self.get_lstm2linear(x)
        outputs = outputs * mask.unsqueeze(-1)
        # 计算bilstm_crf损失
        return -self.crf(outputs, tags, mask.bool())

    # 获取发射分数emission_score
    def get_lstm2linear(self, x):
        embedding = self.word_embeds(x)
        outputs, hidden = self.lstm(embedding)
        outputs = self.dropout(outputs)
        outputs = self.hidden2tag(outputs)
        return outputs

def testNer_bilstm_crf():
    nerlstm_crf = NERLSTM_CRF(config.embedding_dim, config.hidden_dim, config.dropout, word2index, config.tag2id)
    print(f'nerlstm_crf-->\n{nerlstm_crf}')
    # 获取数据
    train_dataloader, dev_dataloader = get_data()
    # 遍历train_dataloader
    for input_ids_padded, labels_padded, attention_mask in train_dataloader:
        result = nerlstm_crf(input_ids_padded, attention_mask)
        print(f'标签预测结果-->\n{result}')
        print('==' * 30)
        loss = nerlstm_crf.log_likelihood(input_ids_padded, labels_padded, attention_mask)
        print(f'loss-->\n{loss}')
        break

if __name__ == '__main__':
    testNer_bilstm_crf()
